# ~*~ coding: utf-8 ~*~
import uuid

from flask_login import current_user
from flask import request, current_app
from flask_socketio import Namespace, join_room

from apps.extensions import db
from apps.assets.models import Asset, SystemUser
from apps.users.models import User
from .models import RequestConnection, WSProxy
from .proxy import ProxyServer
from .utils import get_cached_volume, remove_cache_volume


class SSHSocketIO(Namespace):
    def __init__(self, namespace=None):
        super(SSHSocketIO, self).__init__(namespace=namespace)
        print('init SSHSocketIO ....')

    def _get_current_user(self):
        request.current_user = db.session.query(User).get(current_user.id)

    def _new_connection(self):
        x_forwarded_for = request.headers.get("X-Forwarded-For", '').split(',')
        if x_forwarded_for and x_forwarded_for[0]:
            remote_ip = x_forwarded_for[0]
        else:
            remote_ip = request.remote_addr
        connection = RequestConnection.new_connection(
            addr=(remote_ip, 0), sid=request.sid, sock=self
        )
        connection.user = request.current_user
        connection.login_from = 'WT'

    def on_connect(self):
        print("On connect event trigger")
        self._get_current_user()
        self._new_connection()

    def _connect_host(self, message):
        asset_id = message.get('asset_id', None)
        system_user_id = message.get('system_user_id', None)
        secret = message.get('secret', None)
        width, height = message.get('size', (80, 24))

        connection = RequestConnection.get_connection(request.sid)

        client_id = str(uuid.uuid4())
        client = connection.new_client(client_id)

        client.request.kind = 'session'
        client.request.type = 'pty'
        client.request.meta.update({
            'pty': b'xterm', 'width': width, 'height': height,
        })
        ws_proxy = WSProxy(self, client_id)
        client.chan = ws_proxy
        self.emit('room', {'room': client_id, 'secret': secret})
        join_room(client_id)

        asset = db.session.query(Asset).get(asset_id)
        system_user = db.session.query(SystemUser).get(system_user_id)

        if not asset or not system_user:
            return

        # 使用到数据库创建session对象必须带入应用上下文
        with current_app.app_context():
            forwarder = ProxyServer(client, asset, system_user)

            def proxy():
                try:
                    forwarder.proxy()
                except Exception as e:
                    forwarder.stop()
                    error_message = "Unexpected error occur: {}".format(str(e))
                    print(error_message)
                    client.send(error_message.encode(encoding="utf-8"))
                self.logout(client_id, connection)

            self.socketio.start_background_task(proxy)

    def on_host(self, message):
        # 此处获取主机的信息
        print("On host event trigger")
        self._connect_host(message)

    def on_data(self, message):
        """
        收到浏览器请求
        :param message: {"data": "xxx", "room": "xxx"}
        :return:
        """
        client_id = message.get('room')
        connection = RequestConnection.get_connection(request.sid)
        if not connection:
            return

        client = connection.clients.get(client_id)
        if not client:
            return

        # 将收到的数据发送到服务器
        client.chan.write(message.get("data"))

    def on_disconnect(self):
        print("On disconnect event trigger")
        connection = RequestConnection.get_connection(request.sid)
        if not connection:
            return
        clients_copy = list(connection.clients.keys())
        for client_id in clients_copy:
            try:
                self.logout(client_id, connection)
            except Exception as e:
                print("Unexpected error occur: {}".format(e))
        RequestConnection.remove_connection(connection.id)

    on_close = on_disconnect

    def on_logout(self, client_id):
        print("On logout event trigger")
        connection = RequestConnection.get_connection(request.sid)
        if not connection:
            return
        self.logout(client_id, connection)

    def logout(self, client_id, connection):
        print('logout......')
        connection.remove_client(client_id)
        self.emit('logout', {"room": client_id}, room=client_id)

    def on_ping(self, data):
        print('on_ping-->{}'.format(data))
        self.emit('pong', data)


class SFTPSocketIO(Namespace):
    def on_connect(self):
        self.emit('data', {"sid": str(request.sid)})

    def on_disconnect(self):
        sftp = get_cached_volume(request.sid)
        if sftp:
            sftp.close()
        remove_cache_volume(request.sid)
